"""
Author: xubing
Date: 2024-11-21
Description: 基础LLM客户端类
"""

import os
from abc import ABC, abstractmethod
from typing import Any, Dict, Generator, List, Optional, Union

from pydantic import BaseModel
from myUtils.logger import setup_logger

# 设置日志记录器
logger = setup_logger(__name__)


class Message(BaseModel):
    role: str
    content: Union[str, List[Dict[str, Any]]]


class ChatResponse(BaseModel):
    content: str
    status: bool = True
    error_msg: Optional[str] = None


class BaseClient(ABC):
    """LLM客户端基类"""

    def __init__(self, 
                api_key:str,
                base_url:str,
                model_name: str,
                **kargs):
        self.api_key = api_key
        self.base_url = base_url
        self.model_name = model_name

    @abstractmethod
    def chat(self, messages: List[Dict[str, Any]]) -> ChatResponse:
        """同步聊天接口"""
        pass

    @abstractmethod
    def chat_stream(
            self, messages: List[Dict[str,
                                      Any]]) -> Generator[str, None, None]:
        """流式聊天接口"""
        pass


if __name__ == "__main__":
    print("BaseClient module loaded successfully")
    print(
        f"Model config path: {os.path.join(os.path.dirname(os.path.dirname(__file__)), '_config', 'model_config.yml')}"
    )
